// Copyright 2025 OpenObserve Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use std::sync::Arc;

use bytes::Bytes;
use infra::db::{delete_from_db_coordinator, put_into_db_coordinator};
#[cfg(feature = "enterprise")]
use o2_enterprise::enterprise::common::config::get_config as get_o2_config;

use crate::common::infra::config::USER_SESSIONS;

// Key prefix for session events in coordinator
pub const USER_SESSION_KEY: &str = "/user_sessions/";

pub async fn get(session_id: &str) -> Result<String, anyhow::Error> {
    if let Some(val) = USER_SESSIONS.get(session_id) {
        return Ok(val.to_string());
    }

    // get from db
    log::warn!("Cache miss for user session, read from db: {}", session_id);

    let session = infra::table::sessions::get(session_id).await?;

    match session {
        Some(session) => {
            let access_token = session.access_token.clone();
            // Cache it in memory
            if !access_token.is_empty() {
                USER_SESSIONS.insert(session_id.to_string(), access_token.clone());
            }
            Ok(access_token)
        }
        None => Err(anyhow::anyhow!("Session not found: {}", session_id)),
    }
}

pub async fn set(session_id: &str, val: &str) -> Result<(), anyhow::Error> {
    infra::table::sessions::set(session_id, val).await?;
    let key = format!("{USER_SESSION_KEY}{session_id}");
    if let Err(e) = put_into_db_coordinator(&key, Bytes::new(), true, None).await {
        log::error!("[SESSION] Failed to sync session to coordinator: {key} - {e}");
    }

    #[cfg(feature = "enterprise")]
    if get_o2_config().super_cluster.enabled {
        let _ =
            o2_enterprise::enterprise::super_cluster::queue::put(&key, Bytes::new(), true, None)
                .await
                .inspect_err(|e| {
                    log::error!("[SESSION] put to super cluster failed: {key} - {e}");
                });
    }

    USER_SESSIONS.insert(session_id.to_string(), val.to_string());

    Ok(())
}

pub async fn delete(session_id: &str) -> Result<(), anyhow::Error> {
    infra::table::sessions::delete(session_id).await?;
    let key = format!("{USER_SESSION_KEY}{session_id}");
    if let Err(e) = delete_from_db_coordinator(&key, false, true, None).await {
        log::error!("[SESSION] Failed to delete session from coordinator: {key} - {e}");
    }

    #[cfg(feature = "enterprise")]
    if get_o2_config().super_cluster.enabled {
        let _ = o2_enterprise::enterprise::super_cluster::queue::delete(&key, false, true, None)
            .await
            .inspect_err(|e| {
                log::error!("[SESSION] delete to super cluster failed: {key} - {e}");
            });
    }

    USER_SESSIONS.remove(session_id);

    Ok(())
}

pub async fn watch() -> Result<(), anyhow::Error> {
    let key = USER_SESSION_KEY;
    let cluster_coordinator = infra::db::get_coordinator().await;
    let mut events = cluster_coordinator.watch(key).await?;
    let events = Arc::get_mut(&mut events).unwrap();
    log::info!("Start watching user sessions");
    loop {
        let ev = match events.recv().await {
            Some(ev) => ev,
            None => {
                log::error!("watch_sessions: event channel closed");
                return Ok(());
            }
        };
        match ev {
            infra::db::Event::Put(ev) => {
                let session_id = ev.key.strip_prefix(key).unwrap();
                match infra::table::sessions::get(session_id).await {
                    Ok(Some(session)) => {
                        if !session.access_token.is_empty() {
                            USER_SESSIONS.insert(session_id.to_string(), session.access_token);
                            log::debug!("Session added to cache: {}", session_id);
                        }
                    }
                    Ok(None) => {
                        log::warn!(
                            "Got coordinator message, but session id not found in DB: {}",
                            session_id
                        );
                    }
                    Err(e) => {
                        log::error!(
                            "Error fetching session after coordinator message {}: {}",
                            session_id,
                            e
                        );
                    }
                }
            }
            infra::db::Event::Delete(ev) => {
                let session_id = ev.key.strip_prefix(key).unwrap();
                USER_SESSIONS.remove(session_id);
                log::debug!("Session removed from cache: {}", session_id);
            }
            infra::db::Event::Empty => {}
        }
    }
}

pub async fn cache() -> Result<(), anyhow::Error> {
    let sessions_list = infra::table::sessions::list().await?;

    for session in sessions_list {
        if !session.access_token.is_empty() {
            USER_SESSIONS.insert(session.session_id, session.access_token);
        }
    }

    log::info!(
        "User Sessions Cached: {} sessions loaded",
        USER_SESSIONS.len()
    );
    Ok(())
}
